# Schniepp Lab, 2018-2021

# Atom class to handel atom related calculations

import numpy as np
from Helpers import degree_to_radian, radian_to_degree, rotate_by_degree

class Atom:
    
    # Atom class can be initialized from 3 types of data:
    # 1. from x, y, z, using Cartesian coordinates
    # 2. from r, phi, theta, using spherical coordinates
    # 3. from r_xy, phi, z, using cylindrical coordinates
    def __init__(self, a_type = None, x = None, y = None, z = None, r = None, r_xy = None, \
                 phi = None, theta = None):
        self.a_type = a_type
        self.x = x
        self.y = y
        self.z = z
        self.r = r
        self.r_xy = r_xy
        self.phi = phi
        self.theta = theta
        
        if x != None and y != None and z != None:
            self.r = np.sqrt(x ** 2 + y ** 2 + z ** 2)
            self.r_xy = np.sqrt(x ** 2 + y ** 2)
            
            if z == 0.0:
                self.theta = np.pi / 2
            elif z > 0:
                self.theta = np.pi / 2 - np.arctan(z / self.r_xy)
            else:
                self.theta = np.pi / 2 + np.arctan(-z / self.r_xy)
                
            if x == 0.0:
                if y > 0:
                    self.phi = np.pi / 2
                elif y < 0:
                    self.phi = 3 * np.pi / 2
            elif x > 0:
                if y == 0:
                    self.phi = 0
                elif y > 0:
                    self.phi = np.arctan(y / x)
                elif y < 0:
                    self.phi = 2 * np.pi - np.arctan(-y / x)
            elif x < 0:
                if y == 0:
                    self.phi = np.pi
                elif y > 0:
                    self.phi = np.pi - np.arctan(y / (-x))
                elif y < 0:
                    self.phi = np.pi + np.arctan(y / x)
                    
        elif r != None and phi != None and theta != None:
            if theta < np.pi / 2:
                self.r_xy = r * np.sin(theta)
                self.z = r * np.cos(theta)
            elif theta > np.pi / 2:
                self.r_xy = r * np.sin(np.pi - theta)
                self.z = - r * np.cos(np.pi - theta)
            
            self.x = self.r_xy * np.cos(phi)
            self.y = self.r_xy * np.sin(phi)
            
        elif r_xy != None and phi != None and z != None:
            self.x = self.r_xy * np.cos(phi)
            self.y = self.r_xy * np.sin(phi)
#             print('x = {}, y = {}, z = {}'.format(self.x, self.y, self.z))
            self.r = np.sqrt(self.x ** 2 + self.y ** 2 + self.z ** 2)
            if z == 0:
                self.theta = np.pi / 2
            elif z > 0:
                self.theta = np.pi / 2 - np.arcsin(z / self.r)
            elif z < 0:
                self.theta = np.pi / 2 + np.arcsin(-z / self.r)
                
        if None in [self.a_type, self.x, self.y, self.z, self.r, self.r_xy, self.phi, self.theta]:
            print('Not correctly initialized!')
    
    # return Cartesian coordinates of the atom
    def get_Cartesian(self):
        return (self.x, self.y, self.z)
    
    # return spherical coordinates of the atom
    def get_spherical(self, mode = 'radian'):
        if mode == 'degree':
            return (self.r, radian_to_degree(self.phi), radian_to_degree(self.theta))

        return (self.r, self.phi, self.theta)

    # return cylindrical coordinates of the atom
    def get_cylindrical(self, mode = 'radian'):
        if mode == 'degree':
            return (self.r_xy, radian_to_degree(self.phi), self.z)
        return (self.r_xy, self.phi, self.z)
    
    # print out class information
    def get_info(self):
        print('x = {}, y = {}, z = {}'.format(self.x, self.y, self.z))
        print('r_xy = ', self.r_xy)
        print('r = {}, phi = {}, theta = {}'.format(self.r, self.phi, self.theta))
    
    # return a new Atom instance with the location offset of dz and dd(egree)
    def duplicate_atom_to_location(self, dz, dd):
        return Atom(self.a_type, r_xy = self.r_xy, \
                    phi = degree_to_radian(rotate_by_degree(radian_to_degree(self.phi), dd)), \
                    z = self.z + dz)